{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) 2020, NVIDIA CORPORATION.\n",
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0,1,2,3\"\n",
    "start = time.time()\n",
    "very_start = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.14.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#import pandas as pd, \n",
    "import numpy as np\n",
    "from datetime import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "#pd.set_option('display.max_columns', 500)\n",
    "#pd.set_option('display.max_rows', 500)\n",
    "import cudf, cupy, time, rmm\n",
    "cudf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask as dask, dask_cudf\n",
    "from dask.distributed import Client, wait\n",
    "from dask_cuda import LocalCUDACluster\n",
    "import subprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://10.2.61.36:43723</li>\n",
       "  <li><b>Dashboard: </b><a href='http://10.2.61.36:8787/status' target='_blank'>http://10.2.61.36:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>4</li>\n",
       "  <li><b>Cores: </b>4</li>\n",
       "  <li><b>Memory: </b>270.39 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://10.2.61.36:43723' processes=4 threads=4, memory=270.39 GB>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster = LocalCUDACluster(ip='10.2.61.36')#,protocol=\"ucx\", \n",
    "                           #rmm_pool_size=\"26GB\",\n",
    "                           #enable_tcp_over_ucx=True, enable_nvlink=True)\n",
    "#cluster = LocalCUDACluster()\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.06 s, sys: 352 ms, total: 1.41 s\n",
      "Wall time: 3.48 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>hashtags</th>\n",
       "      <th>tweet_id</th>\n",
       "      <th>links</th>\n",
       "      <th>domains</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>a_follower_count</th>\n",
       "      <th>a_following_count</th>\n",
       "      <th>a_is_verified</th>\n",
       "      <th>a_account_creation</th>\n",
       "      <th>b_follower_count</th>\n",
       "      <th>...</th>\n",
       "      <th>reply</th>\n",
       "      <th>retweet</th>\n",
       "      <th>retweet_comment</th>\n",
       "      <th>like</th>\n",
       "      <th>id</th>\n",
       "      <th>language</th>\n",
       "      <th>tweet_type</th>\n",
       "      <th>media</th>\n",
       "      <th>a_user_id</th>\n",
       "      <th>b_user_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>None</td>\n",
       "      <td>82106226</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581016405</td>\n",
       "      <td>429</td>\n",
       "      <td>605</td>\n",
       "      <td>False</td>\n",
       "      <td>1475647574</td>\n",
       "      <td>65</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>1581018257</td>\n",
       "      <td>0</td>\n",
       "      <td>1581018257</td>\n",
       "      <td>18689</td>\n",
       "      <td>63</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>29278002</td>\n",
       "      <td>17389071</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>None</td>\n",
       "      <td>725094422</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581115466</td>\n",
       "      <td>19125</td>\n",
       "      <td>142</td>\n",
       "      <td>False</td>\n",
       "      <td>1368466630</td>\n",
       "      <td>92</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>18690</td>\n",
       "      <td>59</td>\n",
       "      <td>2</td>\n",
       "      <td>12</td>\n",
       "      <td>21253110</td>\n",
       "      <td>17260406</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>None</td>\n",
       "      <td>187456866</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581369156</td>\n",
       "      <td>554</td>\n",
       "      <td>205</td>\n",
       "      <td>False</td>\n",
       "      <td>1520878712</td>\n",
       "      <td>94</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>1581399261</td>\n",
       "      <td>0</td>\n",
       "      <td>1581399261</td>\n",
       "      <td>18691</td>\n",
       "      <td>59</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>18400787</td>\n",
       "      <td>17260406</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>None</td>\n",
       "      <td>376529151</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581483924</td>\n",
       "      <td>102476</td>\n",
       "      <td>90</td>\n",
       "      <td>False</td>\n",
       "      <td>1507431791</td>\n",
       "      <td>21</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>18692</td>\n",
       "      <td>54</td>\n",
       "      <td>2</td>\n",
       "      <td>4</td>\n",
       "      <td>30346215</td>\n",
       "      <td>99628</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>99AB4BF985AFE868CA7C8458146C5271</td>\n",
       "      <td>356861582</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581391904</td>\n",
       "      <td>1912096</td>\n",
       "      <td>179</td>\n",
       "      <td>True</td>\n",
       "      <td>1245772264</td>\n",
       "      <td>21</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1581414157</td>\n",
       "      <td>18693</td>\n",
       "      <td>54</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>22360189</td>\n",
       "      <td>99628</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 24 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                           hashtags   tweet_id links domains   timestamp  \\\n",
       "0                              None   82106226  None    None  1581016405   \n",
       "1                              None  725094422  None    None  1581115466   \n",
       "2                              None  187456866  None    None  1581369156   \n",
       "3                              None  376529151  None    None  1581483924   \n",
       "4  99AB4BF985AFE868CA7C8458146C5271  356861582  None    None  1581391904   \n",
       "\n",
       "   a_follower_count  a_following_count  a_is_verified  a_account_creation  \\\n",
       "0               429                605          False          1475647574   \n",
       "1             19125                142          False          1368466630   \n",
       "2               554                205          False          1520878712   \n",
       "3            102476                 90          False          1507431791   \n",
       "4           1912096                179           True          1245772264   \n",
       "\n",
       "   b_follower_count  ...  reply     retweet  retweet_comment        like  \\\n",
       "0                65  ...      0  1581018257                0  1581018257   \n",
       "1                92  ...      0           0                0           0   \n",
       "2                94  ...      0  1581399261                0  1581399261   \n",
       "3                21  ...      0           0                0           0   \n",
       "4                21  ...      0           0                0  1581414157   \n",
       "\n",
       "      id  language  tweet_type  media  a_user_id  b_user_id  \n",
       "0  18689        63           1     12   29278002   17389071  \n",
       "1  18690        59           2     12   21253110   17260406  \n",
       "2  18691        59           1     12   18400787   17260406  \n",
       "3  18692        54           2      4   30346215      99628  \n",
       "4  18693        54           1     12   22360189      99628  \n",
       "\n",
       "[5 rows x 24 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "path = '/raid/data/recsys/dask_input/step1_output'\n",
    "df = dask_cudf.read_parquet(f'{path}/*.parquet')\n",
    "df, = dask.persist(df)\n",
    "_ = wait(df)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 299 ms, sys: 14.2 ms, total: 314 ms\n",
      "Wall time: 568 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>hashtags</th>\n",
       "      <th>tweet_id</th>\n",
       "      <th>links</th>\n",
       "      <th>domains</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>a_follower_count</th>\n",
       "      <th>a_following_count</th>\n",
       "      <th>a_is_verified</th>\n",
       "      <th>a_account_creation</th>\n",
       "      <th>b_follower_count</th>\n",
       "      <th>...</th>\n",
       "      <th>like</th>\n",
       "      <th>id</th>\n",
       "      <th>language</th>\n",
       "      <th>tweet_type</th>\n",
       "      <th>media</th>\n",
       "      <th>a_user_id</th>\n",
       "      <th>b_user_id</th>\n",
       "      <th>len_hashtags</th>\n",
       "      <th>len_domains</th>\n",
       "      <th>len_links</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>None</td>\n",
       "      <td>82106226</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581016405</td>\n",
       "      <td>429</td>\n",
       "      <td>605</td>\n",
       "      <td>False</td>\n",
       "      <td>1475647574</td>\n",
       "      <td>65</td>\n",
       "      <td>...</td>\n",
       "      <td>1581018257</td>\n",
       "      <td>18689</td>\n",
       "      <td>63</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>29278002</td>\n",
       "      <td>17389071</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>None</td>\n",
       "      <td>725094422</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581115466</td>\n",
       "      <td>19125</td>\n",
       "      <td>142</td>\n",
       "      <td>False</td>\n",
       "      <td>1368466630</td>\n",
       "      <td>92</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>18690</td>\n",
       "      <td>59</td>\n",
       "      <td>2</td>\n",
       "      <td>12</td>\n",
       "      <td>21253110</td>\n",
       "      <td>17260406</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>None</td>\n",
       "      <td>187456866</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581369156</td>\n",
       "      <td>554</td>\n",
       "      <td>205</td>\n",
       "      <td>False</td>\n",
       "      <td>1520878712</td>\n",
       "      <td>94</td>\n",
       "      <td>...</td>\n",
       "      <td>1581399261</td>\n",
       "      <td>18691</td>\n",
       "      <td>59</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>18400787</td>\n",
       "      <td>17260406</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>None</td>\n",
       "      <td>376529151</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581483924</td>\n",
       "      <td>102476</td>\n",
       "      <td>90</td>\n",
       "      <td>False</td>\n",
       "      <td>1507431791</td>\n",
       "      <td>21</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>18692</td>\n",
       "      <td>54</td>\n",
       "      <td>2</td>\n",
       "      <td>4</td>\n",
       "      <td>30346215</td>\n",
       "      <td>99628</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>99AB4BF985AFE868CA7C8458146C5271</td>\n",
       "      <td>356861582</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>1581391904</td>\n",
       "      <td>1912096</td>\n",
       "      <td>179</td>\n",
       "      <td>True</td>\n",
       "      <td>1245772264</td>\n",
       "      <td>21</td>\n",
       "      <td>...</td>\n",
       "      <td>1581414157</td>\n",
       "      <td>18693</td>\n",
       "      <td>54</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>22360189</td>\n",
       "      <td>99628</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 27 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                           hashtags   tweet_id links domains   timestamp  \\\n",
       "0                              None   82106226  None    None  1581016405   \n",
       "1                              None  725094422  None    None  1581115466   \n",
       "2                              None  187456866  None    None  1581369156   \n",
       "3                              None  376529151  None    None  1581483924   \n",
       "4  99AB4BF985AFE868CA7C8458146C5271  356861582  None    None  1581391904   \n",
       "\n",
       "   a_follower_count  a_following_count  a_is_verified  a_account_creation  \\\n",
       "0               429                605          False          1475647574   \n",
       "1             19125                142          False          1368466630   \n",
       "2               554                205          False          1520878712   \n",
       "3            102476                 90          False          1507431791   \n",
       "4           1912096                179           True          1245772264   \n",
       "\n",
       "   b_follower_count  ...        like     id  language  tweet_type  media  \\\n",
       "0                65  ...  1581018257  18689        63           1     12   \n",
       "1                92  ...           0  18690        59           2     12   \n",
       "2                94  ...  1581399261  18691        59           1     12   \n",
       "3                21  ...           0  18692        54           2      4   \n",
       "4                21  ...  1581414157  18693        54           1     12   \n",
       "\n",
       "   a_user_id  b_user_id  len_hashtags  len_domains  len_links  \n",
       "0   29278002   17389071             0            0          0  \n",
       "1   21253110   17260406             0            0          0  \n",
       "2   18400787   17260406             0            0          0  \n",
       "3   30346215      99628             0            0          0  \n",
       "4   22360189      99628             1            0          0  \n",
       "\n",
       "[5 rows x 27 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "def count_token(ds,token):\n",
    "    not_null = ds.isnull()==0\n",
    "    return ((ds.str.count(token)+1)*not_null).fillna(0)\n",
    "df['len_hashtags'] = df['hashtags'].map_partitions(lambda ds: count_token(ds,'\\t'))\n",
    "df['len_domains']  = df['domains'].map_partitions(lambda ds: count_token(ds,'\\t'))\n",
    "df['len_links']    = df['links'].map_partitions(lambda ds: count_token(ds,'\\t'))\n",
    "\n",
    "df, = dask.persist(df)\n",
    "_ = wait(df)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def most_frequent_token(df, col, idcol, outcol, sep='\\t'):\n",
    "    ds = df[[idcol,col]].compute()\n",
    "    tokens = ds[col].str.tokenize(sep)\n",
    "    \n",
    "    token_counts = ds[col].str.token_count('\\t').to_frame()\n",
    "    token_counts[idcol] = ds[idcol]\n",
    "    del ds\n",
    "    token_rows = reconstruct_token_rows(token_counts,col,idcol)\n",
    "    del token_counts\n",
    "    \n",
    "    tokens = tokens.to_frame()\n",
    "    tokens.columns = ['token']\n",
    "    tokens['token'],_ = tokens['token'].factorize()\n",
    "    tokens['token'] = tokens['token']+1\n",
    "    tokens[idcol] = token_rows\n",
    "    \n",
    "    global_token_counts = tokens['token'].value_counts().reset_index()\n",
    "    global_token_counts.columns = ['token','count']\n",
    "    \n",
    "    tokens = tokens.merge(global_token_counts,on='token',how='left')\n",
    "    del global_token_counts\n",
    "    \n",
    "    token_counts = tokens.groupby(idcol).agg({'count':'max'}).reset_index()\n",
    "    tokens = tokens.merge(token_counts,on=idcol,how='left')\n",
    "    del token_counts\n",
    "    mask = tokens.count_x == tokens.count_y\n",
    "    tokens = tokens[mask]\n",
    "    tokens = tokens.drop(['count_x','count_y'],axis=1)\n",
    "    tokens = tokens.drop_duplicates(subset=['token',idcol])\n",
    "    tokens = tokens.rename(mapper={'token':outcol})\n",
    "    \n",
    "    if outcol == col:\n",
    "        df = df.drop(col,axis=1)\n",
    "    df = df.merge(tokens,on=idcol,how='left')\n",
    "    df[outcol] = df[outcol].fillna(0)\n",
    "    del tokens\n",
    "    return df\n",
    "\n",
    "def reconstruct_token_rows(token_counts,col,idcol):\n",
    "    ids = cupy.asnumpy(token_counts[idcol].values)\n",
    "    rep = cupy.asnumpy(token_counts[col].values)\n",
    "    token_rows = np.repeat(ids,rep)\n",
    "    token_rows = cupy.asarray(token_rows)\n",
    "    return token_rows\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 14.1 s, sys: 16.1 s, total: 30.2 s\n",
      "Wall time: 43 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>tweet_id</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>a_follower_count</th>\n",
       "      <th>a_following_count</th>\n",
       "      <th>a_is_verified</th>\n",
       "      <th>a_account_creation</th>\n",
       "      <th>b_follower_count</th>\n",
       "      <th>b_following_count</th>\n",
       "      <th>b_is_verified</th>\n",
       "      <th>b_account_creation</th>\n",
       "      <th>...</th>\n",
       "      <th>tweet_type</th>\n",
       "      <th>media</th>\n",
       "      <th>a_user_id</th>\n",
       "      <th>b_user_id</th>\n",
       "      <th>len_hashtags</th>\n",
       "      <th>len_domains</th>\n",
       "      <th>len_links</th>\n",
       "      <th>domains</th>\n",
       "      <th>links</th>\n",
       "      <th>hashtags</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>836939972</td>\n",
       "      <td>1580970274</td>\n",
       "      <td>4879</td>\n",
       "      <td>4887</td>\n",
       "      <td>False</td>\n",
       "      <td>1409422288</td>\n",
       "      <td>14675</td>\n",
       "      <td>15049</td>\n",
       "      <td>False</td>\n",
       "      <td>1336421092</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>4</td>\n",
       "      <td>18249810</td>\n",
       "      <td>17151597</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1641851</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>246147565</td>\n",
       "      <td>1581006667</td>\n",
       "      <td>5601</td>\n",
       "      <td>513</td>\n",
       "      <td>False</td>\n",
       "      <td>1236832364</td>\n",
       "      <td>2532</td>\n",
       "      <td>758</td>\n",
       "      <td>False</td>\n",
       "      <td>1237123999</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>12</td>\n",
       "      <td>18235619</td>\n",
       "      <td>17275531</td>\n",
       "      <td>11</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>383668</td>\n",
       "      <td>4411742</td>\n",
       "      <td>1022679</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>17949907</td>\n",
       "      <td>1581530401</td>\n",
       "      <td>10339941</td>\n",
       "      <td>275</td>\n",
       "      <td>True</td>\n",
       "      <td>1175467395</td>\n",
       "      <td>18</td>\n",
       "      <td>197</td>\n",
       "      <td>False</td>\n",
       "      <td>1557021150</td>\n",
       "      <td>...</td>\n",
       "      <td>2</td>\n",
       "      <td>8</td>\n",
       "      <td>17593575</td>\n",
       "      <td>73038</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>86003</td>\n",
       "      <td>5099217</td>\n",
       "      <td>1676966</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>419763933</td>\n",
       "      <td>1581455765</td>\n",
       "      <td>7183600</td>\n",
       "      <td>1334</td>\n",
       "      <td>True</td>\n",
       "      <td>1221538577</td>\n",
       "      <td>456</td>\n",
       "      <td>635</td>\n",
       "      <td>False</td>\n",
       "      <td>1319158411</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>12</td>\n",
       "      <td>29825412</td>\n",
       "      <td>150220</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>316584</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>69281605</td>\n",
       "      <td>1581544301</td>\n",
       "      <td>17484</td>\n",
       "      <td>16983</td>\n",
       "      <td>False</td>\n",
       "      <td>1455267394</td>\n",
       "      <td>22</td>\n",
       "      <td>82</td>\n",
       "      <td>False</td>\n",
       "      <td>1498132616</td>\n",
       "      <td>...</td>\n",
       "      <td>1</td>\n",
       "      <td>8</td>\n",
       "      <td>21671318</td>\n",
       "      <td>71905</td>\n",
       "      <td>3</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1587087</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 27 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "    tweet_id   timestamp  a_follower_count  a_following_count  a_is_verified  \\\n",
       "0  836939972  1580970274              4879               4887          False   \n",
       "1  246147565  1581006667              5601                513          False   \n",
       "2   17949907  1581530401          10339941                275           True   \n",
       "3  419763933  1581455765           7183600               1334           True   \n",
       "4   69281605  1581544301             17484              16983          False   \n",
       "\n",
       "   a_account_creation  b_follower_count  b_following_count  b_is_verified  \\\n",
       "0          1409422288             14675              15049          False   \n",
       "1          1236832364              2532                758          False   \n",
       "2          1175467395                18                197          False   \n",
       "3          1221538577               456                635          False   \n",
       "4          1455267394                22                 82          False   \n",
       "\n",
       "   b_account_creation  ...  tweet_type  media  a_user_id  b_user_id  \\\n",
       "0          1336421092  ...           2      4   18249810   17151597   \n",
       "1          1237123999  ...           2     12   18235619   17275531   \n",
       "2          1557021150  ...           2      8   17593575      73038   \n",
       "3          1319158411  ...           1     12   29825412     150220   \n",
       "4          1498132616  ...           1      8   21671318      71905   \n",
       "\n",
       "   len_hashtags  len_domains  len_links  domains    links  hashtags  \n",
       "0             1            0          0        0        0   1641851  \n",
       "1            11            1          1   383668  4411742   1022679  \n",
       "2             2            1          1    86003  5099217   1676966  \n",
       "3             2            0          0        0        0    316584  \n",
       "4             3            0          0        0        0   1587087  \n",
       "\n",
       "[5 rows x 27 columns]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "for col in ['domains','links','hashtags']:\n",
    "    df = most_frequent_token(df, col=col, idcol='id', outcol=col, sep='\\t')\n",
    "    df, = dask.persist(df)\n",
    "    _ = wait(df)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "121868549 12490390 12490585\n",
      "CPU times: user 1.1 s, sys: 16.2 ms, total: 1.11 s\n",
      "Wall time: 1.46 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "TRAIN_SIZE = 121386431\n",
    "VALID_SIZE = 12434735\n",
    "\n",
    "train = df[df['id']<=TRAIN_SIZE] # df['id'].min() = 1\n",
    "test0 = df[(df['id']>TRAIN_SIZE)&(df['id']<=TRAIN_SIZE+VALID_SIZE)]\n",
    "test1 = df[df['id']>TRAIN_SIZE+VALID_SIZE]\n",
    "\n",
    "print(len(train),len(test0),len(test1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3 µs, sys: 0 ns, total: 3 µs\n",
      "Wall time: 5.72 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "path = '/raid/data/recsys/dask_input/step2_output'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 452 ms, sys: 77.9 ms, total: 530 ms\n",
      "Wall time: 4.54 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "train.to_parquet(f'{path}/train.parquet', write_index=False)\n",
    "test0.to_parquet(f'{path}/valid.parquet', write_index=False)\n",
    "test1.to_parquet(f'{path}/test.parquet', write_index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This notebook took 1.0 minutes\n"
     ]
    }
   ],
   "source": [
    "print('This notebook took %.1f minutes'%((time.time()-very_start)/60.))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
